// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/sparksql/aggregates/register.h>
#include <pollux/functions/sparksql/registration/register.h>

using namespace kumo::pollux::functions::aggregate::test;
using namespace kumo::pollux::plan;

namespace kumo::pollux::functions::aggregate::sparksql::test {
    namespace {
        class CollectSetAggregateTest : public AggregationTestBase {
        protected:
            void SetUp() override {
                AggregationTestBase::SetUp();
                registerAggregateFunctions("");
                functions::sparksql::registerFunctions("spark_");
            }
        };

        TEST_F(CollectSetAggregateTest, global) {
            auto data = make_row_vector({
                make_flat_vector<int32_t>({1, 2, 3, 4, 5, 3, 4, 2, 6, 7}),
            });

            auto expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[1, 2, 3, 4, 5, 6, 7]"}),
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});

            // Null inputs.
            data = make_row_vector({
                make_nullable_flat_vector<int32_t>(
                    {1, 2, std::nullopt, 4, 5, std::nullopt, 4, 2, 6, 7}),
            });

            expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[1, 2, 4, 5, 6, 7]"}),
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});

            // All inputs are null.
            data = make_row_vector({
                make_all_null_flat_vector<int32_t>(1'000),
            });

            expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[]"}),
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});

            // NaN inputs are treated as distinct values.
            data = make_row_vector({
                make_flat_vector<double>(
                    {
                        1,
                        std::numeric_limits<double>::quiet_NaN(),
                        std::nan("1"),
                        std::nan("2")
                    }),
            });

            expected = make_row_vector({
                make_array_vector<double>({
                    {
                        1,
                        std::numeric_limits<double>::quiet_NaN(),
                        std::nan("1"),
                        std::nan("2")
                    },
                }),
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});
        }

        TEST_F(CollectSetAggregateTest, noInputRow) {
            // Empty input.
            auto data = make_row_vector({make_flat_vector<int32_t>({})});
            auto expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[]"}),
            });
            testAggregations({data}, {}, {"collect_set(c0)"}, {}, {expected});

            // No input row passes aggregate filter.
            const auto size = 1'000;
            const auto integers =
                    make_flat_vector<int32_t>(size, [](auto row) { return row % 10; });
            data = make_row_vector({
                integers,
                make_flat_vector<bool>(size, [](auto /*row*/) { return false; }),
            });

            auto op = PlanBuilder()
                    .values({data})
                    .singleAggregation({}, {"collect_set(c0)"}, {"c1"})
                    .planNode();
            assertQuery(op, expected);

            // Some groups have no input row passes aggregate filter.
            data = make_row_vector({
                integers,
                integers,
                make_flat_vector<bool>(size, [](auto row) { return row % 2 == 0; }),
            });
            expected = make_row_vector({
                make_flat_vector<int32_t>({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}),
                make_array_vector_from_json<int32_t>(
                    {"[0]", "[]", "[2]", "[]", "[4]", "[]", "[6]", "[]", "[8]", "[]"}),
            });
            op = PlanBuilder()
                    .values({data})
                    .singleAggregation({"c0"}, {"collect_set(c1)"}, {"c2"})
                    .planNode();
            assertQuery(op, expected);
        }

        TEST_F(CollectSetAggregateTest, groupBy) {
            auto data = make_row_vector({
                make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
                make_flat_vector<int32_t>({1, 2, 3, 4, 5, 3, 4, 2, 6, 7}),
            });

            auto expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector_from_json<int32_t>({"[1, 2, 3, 7]", "[3, 4, 5, 6]"}),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"collect_set(c1)"},
                {"c0", "spark_array_sort(a0)"},
                {expected});

            // Null inputs.
            data = make_row_vector({
                make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
                make_nullable_flat_vector<int32_t>(
                    {
                        1,
                        std::nullopt,
                        3,
                        std::nullopt,
                        5,
                        std::nullopt,
                        3,
                        std::nullopt,
                        6,
                        1
                    }),
            });

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector_from_json<int32_t>({
                    "[1]",
                    "[3, 5, 6]",
                }),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"collect_set(c1)"},
                {"c0", "spark_array_sort(a0)"},
                {expected});

            // All inputs are null for a group.
            data = make_row_vector({
                make_flat_vector<int16_t>({1, 1, 2, 2, 2, 1, 2, 1, 2, 1}),
                make_nullable_flat_vector<int32_t>(
                    {
                        1,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        std::nullopt,
                        1
                    }),
            });

            expected = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_array_vector_from_json<int32_t>({
                    "[1]",
                    "[]",
                }),
            });

            testAggregations(
                {data, data, data},
                {"c0"},
                {"collect_set(c1)"},
                {"c0", "spark_array_sort(a0)"},
                {expected});

            data = make_row_vector({
                make_flat_vector<int16_t>({1, 2}),
                make_map_vector_from_json<int64_t, int64_t>({
                    "{10: 10, 11: 11}",
                    "{12: 12}",
                }),
            });

            testFailingAggregations(
                {data}, {"c0"}, {"collect_set(c1)"}, "Unsupported type MAP");
        }

        TEST_F(CollectSetAggregateTest, arrayWithNestedNulls) {
            auto data = make_row_vector({
                make_array_vector_from_json<int32_t>(
                    {
                        "[1, 2, 3]",
                        "[4 ,5]",
                        "[1, 2, 3]",
                        "[3, 4, 2, 6, 7]",
                        "[1, 2, 3]",
                        "[4 ,5, null]"
                    }),
            });

            auto expected = make_row_vector({
                make_nested_array_vector_from_json<int32_t>(
                    {"[[1, 2, 3], [3, 4, 2, 6, 7], [4 ,5], [4 ,5, null]]"})
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});
        }

        TEST_F(CollectSetAggregateTest, rowWithNestedNull) {
            auto data = make_row_vector({
                make_row_vector({
                    make_flat_vector<int32_t>({1, 2, 2, 1}),
                    make_nullable_flat_vector<int32_t>({1, std::nullopt, 2, 1}),
                })
            });
            auto expected = make_row_vector({
                make_array_of_row_vector(
                    ROW({INTEGER(), INTEGER()}),
                    {
                        {
                            variant::row({1, 1}),
                            variant::row({2, 2}),
                            variant::row({2, variant::null(TypeKind::INTEGER)}),
                        },
                    }),
            });

            testAggregations(
                {data}, {}, {"collect_set(c0)"}, {"spark_array_sort(a0)"}, {expected});
        }

        TEST_F(CollectSetAggregateTest, unknownType) {
            auto data = make_row_vector({
                make_null_constant(TypeKind::UNKNOWN, 3),
            });

            auto expected = make_row_vector({
                make_array_vector_from_json<int32_t>({"[]"}),
            });
            testAggregations({data}, {}, {"collect_set(c0)"}, {}, {expected});

            // The grouping key is of UNKONWN type.
            expected = make_row_vector({
                make_null_constant(TypeKind::UNKNOWN, 1),
                make_array_vector_from_json<int32_t>({"[]"}),
            });
            testAggregations({data}, {"c0"}, {"collect_set(c0)"}, {}, {expected});
        }
    } // namespace
} // namespace kumo::pollux::functions::aggregate::sparksql::test
